#include "Irp.h"
#include "injectk.h"

NTSTATUS IrpGeneral(IN struct _DEVICE_OBJECT* DeviceObject, IN OUT struct _IRP* Irp)
{
    UNREFERENCED_PARAMETER(DeviceObject);
    UNREFERENCED_PARAMETER(Irp);
    Irp->IoStatus.Status = STATUS_SUCCESS;
    Irp->IoStatus.Information = 0;
    IoCompleteRequest(Irp, IO_NO_INCREMENT);
    return STATUS_SUCCESS;
}

NTSTATUS IrpDeviceControl(IN struct _DEVICE_OBJECT* DeviceObject, IN OUT struct _IRP* Irp)
{
    UNREFERENCED_PARAMETER(DeviceObject);

    PIO_STACK_LOCATION lpIoStackLocation = IoGetCurrentIrpStackLocation(Irp);
    ULONG uIoCtlCode = lpIoStackLocation->Parameters.DeviceIoControl.IoControlCode;
    ULONG uInputLength = lpIoStackLocation->Parameters.DeviceIoControl.InputBufferLength;
    PVOID lpBuff = Irp->AssociatedIrp.SystemBuffer;
    NTSTATUS retStatus = STATUS_SUCCESS;
    NTSTATUS IoStatus = STATUS_SUCCESS;
    ULONG uInformation = 0;

    switch (uIoCtlCode)
    {
    case IOCTL_INJECT_TO_PROCESS:
    {
        if (uInputLength == sizeof(IOTYPE_PROCESS_INJECT_INFO))
        {
            PIOTYPE_PROCESS_INJECT_INFO lpInjectInfo = (PIOTYPE_PROCESS_INJECT_INFO)(lpBuff);
            PUNICODE_STRING lpusFileName = NULL;
            SIZE_T stSize;
            BOOLEAN bRet = FALSE;

            do
            {
                if (!lpInjectInfo->Bytes || !lpInjectInfo->lpModuleName) break;

                stSize = lpInjectInfo->Bytes + sizeof(UNICODE_STRING);
                lpusFileName = ExAllocatePool(NonPagedPool, stSize);
                if (!lpusFileName) break;
                RtlZeroMemory(lpusFileName, stSize);
                RtlCopyMemory(lpusFileName + 1, lpInjectInfo->lpModuleName, lpInjectInfo->Bytes);
                lpusFileName->Length = (USHORT)lpInjectInfo->Bytes;
                lpusFileName->MaximumLength = (USHORT)lpInjectInfo->Bytes;
                lpusFileName->Buffer = (PWCH)(lpusFileName + 1);

                //bRet = Mm2InjectDllByInsertApc((HANDLE)(lpInjectInfo->ProcessId), lpusFileName);
                bRet = Mm2InjectDllByHijackThread((HANDLE)(lpInjectInfo->ProcessId), lpusFileName);
                if (!bRet) break;

                bRet = TRUE;

            } while (FALSE);

            if (lpusFileName)
                ExFreePool(lpusFileName);

            if (!bRet)
                retStatus = STATUS_UNSUCCESSFUL;
        }
        else
        {
            retStatus = STATUS_INVALID_PARAMETER;
        }
        break;
    }
    default:
        break;
    }

    Irp->IoStatus.Status = IoStatus;
    Irp->IoStatus.Information = uInformation;
    IoCompleteRequest(Irp, IO_NO_INCREMENT);
    return retStatus;
}
